import asyncio
import json
import base64
import cv2
import numpy as np
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.responses import HTMLResponse
from threading import Thread
import time
import logging
from enum import Enum
from fastapi.templating import Jinja2Templates
import os

# 获取当前文件所在目录，确保路径正确
current_dir = os.path.dirname(os.path.abspath(__file__))
templates_dir = os.path.join(current_dir, "templates")  # 明确指定templates目录位置
templates = Jinja2Templates(directory=templates_dir)

# 设置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 导入功能模块
from vehicle_detection import VehicleDetectionModule
from driver_behavior import DriverBehaviorModule

from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Request, Form
from fastapi.responses import HTMLResponse, RedirectResponse, JSONResponse
from fastapi.templating import Jinja2Templates
from models import User
import uuid
from starlette.middleware.sessions import SessionMiddleware

# 初始化FastAPI应用
app = FastAPI()


# 连接状态枚举
class WebSocketState(Enum):
    CONNECTING = 0
    CONNECTED = 1
    DISCONNECTED = 2


# 初始化用户数据库
user_db = User()

# 添加会话中间件
app.add_middleware(SessionMiddleware, secret_key="your-secret-key-here")  # 生产环境中应使用安全的密钥

# 设置模板目录
templates = Jinja2Templates(directory="templates")


# 新增路由 - 登录页面
@app.get("/login", response_class=HTMLResponse)
async def login_page(request: Request):
    return templates.TemplateResponse("login.html", {"request": request})


# 新增路由 - 注册页面
@app.get("/register", response_class=HTMLResponse)
async def register_page(request: Request):
    return templates.TemplateResponse("register.html", {"request": request})


# 新增路由 - 个人信息页面
@app.get("/profile", response_class=HTMLResponse)
async def profile_page(request: Request):
    if "user" not in request.session:
        return RedirectResponse("/login")
    user_info = user_db.get_user_info(request.session["user"]["user_id"])
    return templates.TemplateResponse("profile.html", {
        "request": request,
        "user": user_info
    })


# 新增路由 - 管理页面
@app.get("/admin", response_class=HTMLResponse)
async def admin_page(request: Request):
    if "user" not in request.session or request.session["user"]["role"] != "admin":
        return RedirectResponse("/login")
    return templates.TemplateResponse("admin.html", {"request": request})


# 处理注册请求
@app.post("/register")
async def register(request: Request, username: str = Form(...),
                   password: str = Form(...), email: str = Form(...)):
    result = user_db.register(username, password, email)
    if result["success"]:
        return RedirectResponse("/login", status_code=303)
    return templates.TemplateResponse("register.html", {
        "request": request,
        "error": result["error"]
    })


# 处理登录请求
@app.post("/login")
async def login(request: Request, username: str = Form(...), password: str = Form(...)):
    result = user_db.login(username, password)
    if result["success"]:
        request.session["user"] = {
            "user_id": result["user_id"],
            "username": result["username"],
            "role": result["role"]
        }
        return RedirectResponse("/", status_code=303)
    return templates.TemplateResponse("login.html", {
        "request": request,
        "error": result["error"]
    })


# 处理登出请求
@app.get("/logout")
async def logout(request: Request):
    request.session.pop("user", None)
    return RedirectResponse("/login")


# 处理更新个人信息请求
@app.post("/update-profile")
async def update_profile(request: Request, email: str = Form(...),
                         password: str = Form(None)):
    if "user" not in request.session:
        return RedirectResponse("/login")

    update_data = {"email": email}
    if password:
        update_data["password"] = password

    result = user_db.update_user_info(request.session["user"]["user_id"], **update_data)
    if result["success"]:
        return RedirectResponse("/profile?success=1")
    return templates.TemplateResponse("profile.html", {
        "request": request,
        "user": user_db.get_user_info(request.session["user"]["user_id"]),
        "error": result["error"]
    })


# 全局状态管理
class ConnectionManager:
    def __init__(self):
        self.active_connections = []
        self.is_running = False
        self.vehicle_module = VehicleDetectionModule()
        self.behavior_module = DriverBehaviorModule()
        self.vehicle_camera = None
        self.driver_camera = None
        self.distance_threshold = 30
        self.frame_threshold = 15
        self.vehicle_camera_index = 0
        self.driver_camera_index = 0
        self.camera_resolution = (640, 480)
        self.use_single_camera = False
        self.frame_counter = 0  # 用于调试的帧计数器
        self.last_vehicle_frame = None  # 缓存最后一帧车辆画面
        self.last_driver_frame = None  # 缓存最后一帧驾驶员画面

    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        self.active_connections.append(websocket)
        self.behavior_module.CONSECUTIVE_FRAMES = self.frame_threshold
        self.vehicle_module.warning_threshold = self.distance_threshold

    def disconnect(self, websocket: WebSocket):
        if websocket in self.active_connections:
            self.active_connections.remove(websocket)

    async def broadcast(self, message: dict):
        """向所有连接的客户端广播消息"""
        for connection in self.active_connections:
            try:
                await connection.send_text(json.dumps(message))
            except (WebSocketDisconnect, RuntimeError) as e:
                logger.error(f"广播失败: {str(e)}")
                self.disconnect(connection)
            except Exception as e:
                logger.error(f"广播异常: {str(e)}")
                self.disconnect(connection)

    async def send_log(self, message, is_error=False):
        """发送日志消息"""
        log_level = "ERROR" if is_error else "INFO"
        logger.log(logging.ERROR if is_error else logging.INFO, message)
        await self.broadcast({
            "type": "log",
            "message": message,
            "is_error": is_error
        })

    def release_cameras(self):
        """释放摄像头资源"""
        if self.vehicle_camera is not None:
            self.vehicle_camera.release()
            self.vehicle_camera = None
        if self.driver_camera is not None and not self.use_single_camera:
            self.driver_camera.release()
            self.driver_camera = None
        elif self.driver_camera is not None and self.use_single_camera:
            self.driver_camera = None  # 避免重复释放
        self.last_vehicle_frame = None
        self.last_driver_frame = None

    def find_available_cameras(self):
        """检测可用的摄像头索引"""
        available_indexes = []
        for index in range(0, 4):
            cap = cv2.VideoCapture(index)
            if cap.isOpened():
                # 设置低分辨率提高检测速度
                cap.set(cv2.CAP_PROP_FRAME_WIDTH, 320)
                cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 240)
                ret, frame = cap.read()
                if ret and frame is not None:
                    available_indexes.append(index)
                cap.release()
        return available_indexes

    async def start_cameras(self):
        """启动摄像头"""
        self.release_cameras()
        self.frame_counter = 0  # 重置帧计数器

        try:
            available_cameras = self.find_available_cameras()
            await self.send_log(f"检测到可用摄像头索引: {available_cameras}", False)

            if not available_cameras:
                raise Exception("没有检测到可用的摄像头")

            # 如果只有一个摄像头可用，使用同一个摄像头
            if len(available_cameras) == 1:
                self.use_single_camera = True
                self.vehicle_camera_index = available_cameras[0]
                self.driver_camera_index = available_cameras[0]
                await self.send_log("检测到只有一个摄像头，将使用同一个摄像头进行车辆检测和驾驶员监控", False)
            else:
                self.use_single_camera = False

                # 确保两个摄像头使用不同的索引
                if self.vehicle_camera_index == self.driver_camera_index:
                    # 如果当前设置相同，自动分配不同的摄像头
                    if available_cameras[0] == self.vehicle_camera_index:
                        self.driver_camera_index = available_cameras[1]
                    else:
                        self.vehicle_camera_index = available_cameras[0]
                        self.driver_camera_index = available_cameras[1]

                # 验证索引是否在可用列表中
                if self.vehicle_camera_index not in available_cameras:
                    self.vehicle_camera_index = available_cameras[0]

                if self.driver_camera_index not in available_cameras:
                    self.driver_camera_index = available_cameras[1] if len(available_cameras) > 1 else \
                        available_cameras[0]

            await self.send_log(f"尝试使用车辆摄像头索引: {self.vehicle_camera_index}", False)
            await self.send_log(f"尝试使用驾驶员摄像头索引: {self.driver_camera_index}", False)

            # 尝试打开车辆摄像头
            self.vehicle_camera = cv2.VideoCapture(self.vehicle_camera_index)
            if not self.vehicle_camera.isOpened():
                # 尝试其他可用摄像头
                for index in available_cameras:
                    if index != self.driver_camera_index or self.use_single_camera:
                        self.vehicle_camera = cv2.VideoCapture(index)
                        if self.vehicle_camera.isOpened():
                            self.vehicle_camera_index = index
                            await self.send_log(f"使用索引 {index} 作为车辆摄像头", False)
                            break
                if not self.vehicle_camera.isOpened():
                    raise Exception(f"无法打开车辆摄像头（尝试了索引 {available_cameras}）")

            # 设置摄像头参数 - 关键优化点
            self.vehicle_camera.set(cv2.CAP_PROP_FRAME_WIDTH, self.camera_resolution[0])
            self.vehicle_camera.set(cv2.CAP_PROP_FRAME_HEIGHT, self.camera_resolution[1])
            self.vehicle_camera.set(cv2.CAP_PROP_FPS, 30)
            self.vehicle_camera.set(cv2.CAP_PROP_BUFFERSIZE, 1)  # 减少缓冲区大小
            self.vehicle_camera.set(cv2.CAP_PROP_AUTOFOCUS, 0)  # 关闭自动对焦

            # 测试车辆摄像头
            for _ in range(5):  # 清空缓冲区
                self.vehicle_camera.grab()
            ret1, frame1 = self.vehicle_camera.retrieve()
            if not ret1 or frame1 is None:
                self.release_cameras()
                raise Exception(f"车辆摄像头（索引 {self.vehicle_camera_index}）已打开，但无法读取画面")

            # 保存车辆测试帧
            self.last_vehicle_frame = frame1.copy()
            cv2.imwrite("vehicle_cam_test_frame.jpg", frame1)

            # 尝试打开驾驶员摄像头
            if not self.use_single_camera and self.driver_camera_index != self.vehicle_camera_index:
                try:
                    await self.send_log(f"尝试打开独立的驾驶员摄像头（索引 {self.driver_camera_index})", False)
                    self.driver_camera = cv2.VideoCapture(self.driver_camera_index)
                    if not self.driver_camera.isOpened():
                        raise Exception(f"无法打开驾驶员摄像头索引 {self.driver_camera_index}")

                    # 设置驾驶员摄像头参数
                    self.driver_camera.set(cv2.CAP_PROP_FRAME_WIDTH, self.camera_resolution[0])
                    self.driver_camera.set(cv2.CAP_PROP_FRAME_HEIGHT, self.camera_resolution[1])
                    self.driver_camera.set(cv2.CAP_PROP_FPS, 30)
                    self.driver_camera.set(cv2.CAP_PROP_BUFFERSIZE, 1)
                    self.driver_camera.set(cv2.CAP_PROP_AUTOFOCUS, 0)

                    # 测试驾驶员摄像头
                    for _ in range(5):  # 清空缓冲区
                        self.driver_camera.grab()
                    ret2, frame2 = self.driver_camera.retrieve()
                    if not ret2 or frame2 is None:
                        raise Exception(f"驾驶员摄像头（索引 {self.driver_camera_index}）已打开，但无法读取画面")

                    # 保存驾驶员测试帧
                    self.last_driver_frame = frame2.copy()
                    cv2.imwrite("driver_cam_test_frame.jpg", frame2)
                    await self.send_log(f"驾驶员摄像头（索引 {self.driver_camera_index}）初始化成功", False)
                except Exception as e:
                    await self.send_log(f"独立驾驶员摄像头错误: {str(e)}，将回退到单摄像头模式", True)
                    if self.driver_camera is not None:
                        self.driver_camera.release()
                    self.use_single_camera = True
                    self.driver_camera = self.vehicle_camera
                    self.driver_camera_index = self.vehicle_camera_index
            else:
                # 使用同一个摄像头
                self.use_single_camera = True
                self.driver_camera = self.vehicle_camera
                self.driver_camera_index = self.vehicle_camera_index
                await self.send_log("使用同一个摄像头进行车辆检测和驾驶员监控", False)

            await self.send_log(
                f"摄像头初始化成功! 车辆摄像头: 索引 {self.vehicle_camera_index}, 驾驶员摄像头: 索引 {self.driver_camera_index}",
                False)

            if self.use_single_camera:
                await self.send_log("注意: 使用同一个摄像头进行车辆检测和驾驶员监控，可能会影响性能", True)

            return True
        except Exception as e:
            await self.send_log(f"摄像头初始化失败: {str(e)}", True)
            return False

    def process_frames(self):
        """处理视频帧并发送给前端"""
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)

        try:
            while self.is_running and len(self.active_connections) > 0:
                # 处理车辆检测
                vehicle_frame = None
                if self.vehicle_camera is not None and self.vehicle_camera.isOpened():
                    # 清空缓冲区 - 关键优化点
                    for _ in range(2):
                        self.vehicle_camera.grab()

                    ret1, vehicle_frame = self.vehicle_camera.retrieve()

                    # 更新缓存帧
                    if ret1 and vehicle_frame is not None:
                        self.last_vehicle_frame = vehicle_frame.copy()
                    else:
                        # 如果当前帧读取失败，使用缓存的最后一帧
                        if self.last_vehicle_frame is not None:
                            vehicle_frame = self.last_vehicle_frame.copy()
                            ret1 = True  # 标记为有效帧

                # 如果没有获取到新的帧，使用缓存的最后一帧
                if vehicle_frame is None and self.last_vehicle_frame is not None:
                    vehicle_frame = self.last_vehicle_frame.copy()
                    ret1 = True  # 标记为有效帧

                if vehicle_frame is not None and ret1:
                    # 调试：每100帧保存一次图像
                    self.frame_counter += 1
                    if self.frame_counter % 100 == 0:
                        cv2.imwrite(f"debug_vehicle_{self.frame_counter}.jpg", vehicle_frame)

                    try:
                        vehicle_result = self.vehicle_module.process_frame(vehicle_frame)
                    except Exception as e:
                        loop.run_until_complete(self.send_log(f"车辆检测处理错误: {str(e)}", True))
                        continue

                    # 将帧转换为base64编码
                    try:
                        # 添加图像有效性检查
                        if vehicle_result["frame"] is None or vehicle_result["frame"].size == 0:
                            raise ValueError("车辆检测返回空帧")

                        _, buffer = cv2.imencode('.jpg', vehicle_result["frame"],
                                                 [int(cv2.IMWRITE_JPEG_QUALITY), 80])
                        if buffer is None:
                            raise ValueError("图像编码失败")

                        frame_base64 = base64.b64encode(buffer).decode('utf-8')
                    except Exception as e:
                        loop.run_until_complete(self.send_log(f"车辆帧编码错误: {str(e)}", True))
                        continue

                    # 发送车辆检测结果
                    loop.run_until_complete(self.broadcast({
                        "type": "vehicle",
                        "vehicle_count": vehicle_result["vehicle_count"],
                        "closest_distance": vehicle_result["closest_distance"],
                        "warning": vehicle_result["warning"],
                        "frame": frame_base64
                    }))

                # 处理驾驶员行为识别
                driver_frame = None
                if not self.use_single_camera and self.driver_camera is not None and self.driver_camera.isOpened():
                    # 清空缓冲区
                    for _ in range(2):
                        self.driver_camera.grab()
                    ret2, driver_frame = self.driver_camera.retrieve()

                    # 更新缓存帧
                    if ret2 and driver_frame is not None:
                        self.last_driver_frame = driver_frame.copy()
                    else:
                        # 如果当前帧读取失败，使用缓存的最后一帧
                        if self.last_driver_frame is not None:
                            driver_frame = self.last_driver_frame.copy()
                            ret2 = True  # 标记为有效帧
                else:
                    # 使用车辆摄像头或缓存帧
                    if vehicle_frame is not None:
                        driver_frame = vehicle_frame.copy()
                        self.last_driver_frame = driver_frame.copy()
                        ret2 = True
                    elif self.last_driver_frame is not None:
                        driver_frame = self.last_driver_frame.copy()
                        ret2 = True

                # 如果没有获取到新的帧，使用缓存的最后一帧
                if driver_frame is None and self.last_driver_frame is not None:
                    driver_frame = self.last_driver_frame.copy()
                    ret2 = True

                if driver_frame is not None and ret2:
                    # 调试
                    if self.frame_counter % 100 == 0:
                        cv2.imwrite(f"debug_driver_{self.frame_counter}.jpg", driver_frame)

                    try:
                        # 镜像翻转
                        driver_frame = cv2.flip(driver_frame, 1)
                        behavior_result = self.behavior_module.process_frame(driver_frame)
                    except Exception as e:
                        loop.run_until_complete(self.send_log(f"驾驶员行为检测错误: {str(e)}", True))
                        continue

                    # 将帧转换为base64编码
                    try:
                        if behavior_result["frame"] is None or behavior_result["frame"].size == 0:
                            raise ValueError("行为检测返回空帧")

                        _, buffer = cv2.imencode('.jpg', behavior_result["frame"],
                                                 [int(cv2.IMWRITE_JPEG_QUALITY), 80])
                        if buffer is None:
                            raise ValueError("驾驶员图像编码失败")

                        frame_base64 = base64.b64encode(buffer).decode('utf-8')
                    except Exception as e:
                        loop.run_until_complete(self.send_log(f"驾驶员帧编码错误: {str(e)}", True))
                        continue

                    # 发送行为识别结果
                    loop.run_until_complete(self.broadcast({
                        "type": "driver",
                        "behaviors": behavior_result["behaviors"],
                        "frame": frame_base64
                    }))

                # 控制帧率
                time.sleep(0.03)  # 约33fps
        except Exception as e:
            loop.run_until_complete(self.send_log(f"处理帧时出错: {str(e)}", True))
        finally:
            self.release_cameras()
            self.is_running = False
            loop.close()

    def start_processing(self):
        """开始处理视频流"""
        if not self.is_running:
            self.is_running = True
            thread = Thread(target=self.process_frames, daemon=True)
            thread.start()
            return True
        return False

    def stop_processing(self):
        """停止处理视频流"""
        self.is_running = False
        return True

    def update_parameters(self, params):
        """更新系统参数"""
        if "distance_threshold" in params:
            self.distance_threshold = params["distance_threshold"]
            self.vehicle_module.warning_threshold = self.distance_threshold

        if "frame_threshold" in params:
            self.frame_threshold = params["frame_threshold"]
            self.behavior_module.CONSECUTIVE_FRAMES = self.frame_threshold

        if "vehicle_camera_index" in params:
            self.vehicle_camera_index = int(params["vehicle_camera_index"])

        if "driver_camera_index" in params:
            self.driver_camera_index = int(params["driver_camera_index"])

        return True


# 创建连接管理器实例
manager = ConnectionManager()


# 提供前端页面
@app.get("/", response_class=HTMLResponse)
async def get_index(request: Request):
    if "user" not in request.session:
        return RedirectResponse("/login")

    # 使用模板引擎渲染，确保模板变量正确解析
    return templates.TemplateResponse("index1.html", {"request": request})

# WebSocket端点
@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    # 检查用户是否登录
    cookies = websocket.cookies  # 使用公共API获取cookies
    session_id = cookies.get("session")

    # 这里简化处理，实际应用中需要验证session_id有效性
    # 生产环境中应使用更安全的会话验证机制
    if not session_id:
        await websocket.close(code=1008, reason="未授权访问")
        return
    await manager.connect(websocket)
    try:
        while True:
            data = await websocket.receive_text()
            try:
                message = json.loads(data)
            except json.JSONDecodeError:
                await manager.send_log("无法解析的JSON数据", True)
                continue
            except WebSocketDisconnect:
                manager.disconnect(websocket)
                await manager.send_log("客户端已断开连接")
            except Exception as e:
                await manager.send_log(f"WebSocket错误: {str(e)}", True)
                manager.disconnect(websocket)
                break  # 发生异常后退出循环

            if message["command"] == "start":
                if "params" in message:
                    manager.update_parameters(message["params"])
                    if "vehicle_camera_index" in message["params"]:
                        await manager.send_log(f"车辆摄像头索引更新为: {manager.vehicle_camera_index}", False)
                    if "driver_camera_index" in message["params"]:
                        await manager.send_log(f"驾驶员摄像头索引更新为: {manager.driver_camera_index}", False)

                if await manager.start_cameras():
                    await manager.send_log("摄像头启动成功")
                    if manager.start_processing():
                        await manager.send_log("开始处理视频流")
                    else:
                        await manager.send_log("无法启动视频处理", True)
                else:
                    await manager.send_log("启动摄像头失败", True)

            elif message["command"] == "stop":
                if manager.stop_processing():
                    await manager.send_log("已停止视频处理")
                manager.release_cameras()

            elif message["command"] == "reset":
                manager.stop_processing()
                await manager.send_log("正在重置系统...")
                manager.vehicle_module = VehicleDetectionModule()
                manager.vehicle_module.warning_threshold = manager.distance_threshold
                manager.behavior_module = DriverBehaviorModule()
                manager.behavior_module.CONSECUTIVE_FRAMES = manager.frame_threshold
                manager.use_single_camera = False

                if await manager.start_cameras():
                    if manager.start_processing():
                        await manager.send_log("系统已重置并重新启动")
                    else:
                        await manager.send_log("重置后无法启动视频处理", True)
                else:
                    await manager.send_log("重置后启动摄像头失败", True)

            elif message["command"] == "detect_cameras":
                available_cameras = manager.find_available_cameras()
                await websocket.send_text(json.dumps({
                    "type": "camera_list",
                    "available_cameras": available_cameras
                }))

    except WebSocketDisconnect:
        manager.disconnect(websocket)
        await manager.send_log("客户端已断开连接")
    except Exception as e:
        await manager.send_log(f"WebSocket错误: {str(e)}", True)
        manager.disconnect(websocket)


if __name__ == "__main__":
    import uvicorn

    uvicorn.run("websocket_server:app", host="0.0.0.0", port=8000, reload=True)